package com.tom.learnbase.two;

import cn.hutool.core.io.resource.ResourceUtil;
import com.tom.contants.Contants;
import dev.langchain4j.chain.ConversationalRetrievalChain;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.embedding.AllMiniLmL6V2EmbeddingModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.input.Prompt;
import dev.langchain4j.model.input.PromptTemplate;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiModelName;
import dev.langchain4j.model.openai.OpenAiTokenizer;
import dev.langchain4j.retriever.EmbeddingStoreRetriever;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import org.junit.jupiter.api.Test;

import java.net.URISyntaxException;
import java.net.URL;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.time.Duration;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * 根据文档对话
 *
 */
public class ChatWithDocumentsExamples {


    /**
     * 分词
     */
    @Test
    public void DocumentSplittersTest(){
        /**
         * 文件中
         * txt文档
         */
        Document document = FileSystemDocumentLoader.loadDocument(toPath("example-files/story-about-happy-carrot.txt"), new TextDocumentParser());
        System.out.println(document.text());
        DocumentSplitter splitter = DocumentSplitters.recursive(300, 0);
        List<TextSegment> split = splitter.split(document);
        System.out.println("-----------------");
        split.stream().forEach(textSegment -> System.out.println(textSegment.text()) );
    }


    /**
     * 分词
     */
    @Test
    public void DocumentSplitters2Test(){
        /**
         * 文件中
         * txt文档
         */
        Document document = FileSystemDocumentLoader.loadDocument(toPath("example-files/story-about-happy-carrot.txt"), new TextDocumentParser());
        System.out.println(document.text());
        DocumentSplitter splitter = DocumentSplitters.recursive(
                100,
                0,
                new OpenAiTokenizer(OpenAiModelName.GPT_3_5_TURBO)
        );
        List<TextSegment> split = splitter.split(document);
        System.out.println("-----------------");
        split.stream().forEach(textSegment -> System.out.println(textSegment.text()) );
    }

    /**
     * 向量化
     */
    @Test
    public void TestEmbedd(){

        EmbeddingModel embeddingModel = new AllMiniLmL6V2EmbeddingModel();

        /**
         * 内存级别的存储
         * 底层是CopyOnWriteArrayList
         */
        EmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();

        /**
         * DocumentSplitters 分词器
         */
        EmbeddingStoreIngestor ingestor = EmbeddingStoreIngestor.builder()
                .documentSplitter(DocumentSplitters.recursive(300, 0))
                .embeddingModel(embeddingModel)
                .embeddingStore(embeddingStore)
                .build();

        /**
         * 文件中
         * txt文档
         */
        Document document = FileSystemDocumentLoader.loadDocument(toPath("example-files/story-about-happy-carrot.txt"), new TextDocumentParser());
        /**
         * 文档内容格式化
         * 文档内容分词
         * 词语向量化，并存储
         */
        ingestor.ingest(document);


        EmbeddingStoreRetriever storeRetriever = EmbeddingStoreRetriever.from(embeddingStore, embeddingModel);
        //这里没有对 "Who is Charlie?"进行分词，然后向量化，所以不准确
        List<TextSegment> relevant = storeRetriever.findRelevant("Who is Charlie?");
        String infor=(String)relevant.stream().map(TextSegment::text).map((segment) -> {
            return "..." + segment + "...";
        }).collect(Collectors.joining("\n\n"));
        System.out.println(infor);

    }
















    // Please also check ServiceWithRetrieverExample

    static class IfYouNeedSimplicity {

        public static void main(String[] args) throws Exception {

            EmbeddingModel embeddingModel = new AllMiniLmL6V2EmbeddingModel();

            /**
             * 内存级别的存储
             * 底层是CopyOnWriteArrayList
             */
            EmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();

            /**
             * DocumentSplitters 分词器
             */
            EmbeddingStoreIngestor ingestor = EmbeddingStoreIngestor.builder()
                    .documentSplitter(DocumentSplitters.recursive(300, 0))
                    .embeddingModel(embeddingModel)
                    .embeddingStore(embeddingStore)
                    .build();

            /**
             * 文件中
             * txt文档
             */
            Document document = FileSystemDocumentLoader.loadDocument(toPath("example-files/story-about-happy-carrot.txt"), new TextDocumentParser());
            /**
             * 文档内容格式化
             * 文档内容分词
             * 词语向量化，并存储
             */
            ingestor.ingest(document);




            ConversationalRetrievalChain chain = ConversationalRetrievalChain.builder()
                    .chatLanguageModel(OpenAiChatModel.withApiKey(Contants.key))
                    .retriever(EmbeddingStoreRetriever.from(embeddingStore, embeddingModel))
                    // .chatMemory() // you can override default chat memory
                    // .promptTemplate() // you can override default prompt template
                    .build();


            /**
             * 把问题 向量化 A
             * 从存储向量遍历每个向量 B 计算两个嵌入向量之间的余弦相似度 cosineSimilarity
             * (cosineSimilarity + 1.0) / 2.0 将余弦相似度转换为范围在0到1之间的打分
             * 如果高于最低打分，则该向量挑选出来，组合成list，并根据打分排序 ，list组合成information
             * 使用PromptTemplate 形成Prompt
             * PromptTemplate.from("Answer the following question to the best of your ability: {{question}}\n\nBase your answer on the following information:\n{{information}}")
             *
             * 访问llms
             *
             *
             */
            String answer = chain.execute("Who is Charlie?");
            System.out.println(answer); // Charlie is a cheerful carrot living in VeggieVille...
        }
    }

    static class If_You_Need_More_Control {

        public static void main(String[] args) {

            // Load the document that includes the information you'd like to "chat" about with the model.
            Document document = FileSystemDocumentLoader.loadDocument(toPath("example-files/story-about-happy-carrot.txt"), new TextDocumentParser());

            // Split document into segments 100 tokens each
            DocumentSplitter splitter = DocumentSplitters.recursive(
                    100,
                    0,
                    new OpenAiTokenizer(OpenAiModelName.GPT_3_5_TURBO)
            );
            List<TextSegment> segments = splitter.split(document);

            // Embed segments (convert them into vectors that represent the meaning) using embedding model
            EmbeddingModel embeddingModel = new AllMiniLmL6V2EmbeddingModel();
            List<Embedding> embeddings = embeddingModel.embedAll(segments).content();

            // Store embeddings into embedding store for further search / retrieval
            EmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();
            embeddingStore.addAll(embeddings, segments);

            // Specify the question you want to ask the model
            String question = "Who is Charlie?";

            // Embed the question
            Embedding questionEmbedding = embeddingModel.embed(question).content();

            // Find relevant embeddings in embedding store by semantic similarity
            // You can play with parameters below to find a sweet spot for your specific use case
            int maxResults = 3;
            double minScore = 0.7;
            List<EmbeddingMatch<TextSegment>> relevantEmbeddings
                    = embeddingStore.findRelevant(questionEmbedding, maxResults, minScore);

            // Create a prompt for the model that includes question and relevant embeddings
            PromptTemplate promptTemplate = PromptTemplate.from(
                    "Answer the following question to the best of your ability:\n"
                            + "\n"
                            + "Question:\n"
                            + "{{question}}\n"
                            + "\n"
                            + "Base your answer on the following information:\n"
                            + "{{information}}");

            String information = relevantEmbeddings.stream()
                    .map(match -> match.embedded().text())
                    .collect(Collectors.joining("\n\n"));
            /**
             * Charlie's idea had saved the day, but he humbly attributed the success to their teamwork and friendship.
             * They celebrated their victory with a grand party, filled with laughter, dance, and merry games.
             * That night, under the twinkling stars, they made a pact to always stand by each other, come what may.
             *
             * Once upon a time in the town of VeggieVille, there lived a cheerful carrot named Charlie.
             * Charlie was a radiant carrot, always beaming with joy and positivity.
             * His vibrant orange skin and lush green top were a sight to behold, but it was his infectious laughter and warm personality that really set him apart.
             *
             * Charlie had a diverse group of friends, each a vegetable with their own unique characteristics.
             * There was Bella the blushing beetroot, always ready with a riddle or two; Timmy the timid tomato, a gentle soul with a heart of gold; and Percy the prankster potato, whose jokes always brought a smile to everyone's faces.
             * Despite their differences, they shared a close bond, their friendship as robust as their natural goodness.
             * Charlie is a cheerful and radiant carrot who lived in the town of VeggieVille. He had a diverse group of friends, including Bella the beetroot, Timmy the tomato, and Percy the potato. Charlie's idea saved the day for them all, and he celebrated their victory with a grand party. Despite his success, Charlie humbly attributed their teamwork and friendship to their success. They made a pact to always stand by each other, showing the depth of their bond and loyalty.
             *
             */
            System.out.println("information:"+information);
            Map<String, Object> variables = new HashMap<>();
            variables.put("question", question);
            variables.put("information", information);

            Prompt prompt = promptTemplate.apply(variables);

            // Send the prompt to the OpenAI chat model
            ChatLanguageModel chatModel = OpenAiChatModel.builder()
                    .apiKey(Contants.key)
                    .timeout(Duration.ofSeconds(60))
                    .build();
            AiMessage aiMessage = chatModel.generate(prompt.toUserMessage()).content();

            // See an answer from the model  准确很多
            String answer = aiMessage.text();
            System.out.println(answer); // Charlie is a cheerful carrot living in VeggieVille...
        }
    }

    private static Path toPath(String fileName) {
        try {
            URL fileUrl = ResourceUtil.getResource(fileName);
            return Paths.get(fileUrl.toURI());
        } catch (URISyntaxException e) {
            throw new RuntimeException(e);
        }
    }
}
